"""Neo4j客户端 - 关联记忆存储"""

import asyncio
from typing import Dict, Any, Optional, List, Union, Tuple
from datetime import datetime
from contextlib import asynccontextmanager
from loguru import logger

try:
    from neo4j import AsyncGraphDatabase, AsyncDriver, AsyncSession
    from neo4j.exceptions import ServiceUnavailable, AuthError, TransientError, ClientError
except ImportError:
    raise ImportError("请安装neo4j: pip install neo4j")

from src.config import get_database_url


class Neo4jClient:
    """Neo4j客户端 - 负责关联记忆存储"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.neo4j_config = config.get('databases', {}).get('neo4j', {})
        self.driver: Optional[AsyncDriver] = None
        self._connected = False
        self._max_retry_attempts = 3
        self._retry_delay = 1.0  # 秒
    
    async def connect(self):
        """连接Neo4j"""
        if self._connected:
            return
        
        uri = self.neo4j_config.get('uri', 'bolt://localhost:7687')
        username = self.neo4j_config.get('username', 'neo4j')
        password = self.neo4j_config.get('password', 'password')
        
        for attempt in range(self._max_retry_attempts):
            try:
                self.driver = AsyncGraphDatabase.driver(
                    uri,
                    auth=(username, password),
                    max_connection_lifetime=self.neo4j_config.get('max_connection_lifetime', 3600),
                    max_connection_pool_size=self.neo4j_config.get('max_pool_size', 50),
                    connection_acquisition_timeout=self.neo4j_config.get('connection_timeout', 60)
                )
                
                # 测试连接
                await self.driver.verify_connectivity()
                
                # 创建约束和索引
                await self._create_constraints_and_indexes()
                
                self._connected = True
                logger.info("Neo4j连接成功")
                return
                
            except (ServiceUnavailable, TransientError) as e:
                if attempt < self._max_retry_attempts - 1:
                    logger.warning(f"Neo4j连接失败，尝试重连 ({attempt + 1}/{self._max_retry_attempts}): {e}")
                    await asyncio.sleep(self._retry_delay * (2 ** attempt))  # 指数退避
                    continue
                else:
                    logger.error(f"Neo4j连接失败，已达到最大重试次数: {e}")
                    raise
            except (AuthError, ClientError) as e:
                logger.error(f"Neo4j连接失败: {e}")
                raise
            except Exception as e:
                logger.error(f"Neo4j连接异常: {e}")
                raise
    
    async def close(self):
        """关闭Neo4j连接"""
        if self.driver:
            await self.driver.close()
            self.driver = None
            self._connected = False
            logger.info("Neo4j连接已关闭")
    
    async def ping(self) -> bool:
        """检查连接状态"""
        try:
            if self.driver:
                await self.driver.verify_connectivity()
                return True
            return False
        except Exception:
            return False
    
    async def _create_constraints_and_indexes(self):
        """创建约束和索引"""
        try:
            async with self.driver.session() as session:
                # 创建唯一约束
                constraints = [
                    "CREATE CONSTRAINT user_id_unique IF NOT EXISTS FOR (u:User) REQUIRE u.user_id IS UNIQUE",
                    "CREATE CONSTRAINT concept_id_unique IF NOT EXISTS FOR (c:Concept) REQUIRE c.concept_id IS UNIQUE",
                    "CREATE CONSTRAINT memory_id_unique IF NOT EXISTS FOR (m:Memory) REQUIRE m.memory_id IS UNIQUE",
                    "CREATE CONSTRAINT event_id_unique IF NOT EXISTS FOR (e:Event) REQUIRE e.event_id IS UNIQUE"
                ]
                
                for constraint in constraints:
                    try:
                        await session.run(constraint)
                    except Exception as e:
                        # 约束可能已存在，忽略错误
                        logger.debug(f"约束创建跳过: {e}")
                
                # 创建索引
                indexes = [
                    "CREATE INDEX user_name_index IF NOT EXISTS FOR (u:User) ON (u.name)",
                    "CREATE INDEX concept_name_index IF NOT EXISTS FOR (c:Concept) ON (c.name)",
                    "CREATE INDEX memory_timestamp_index IF NOT EXISTS FOR (m:Memory) ON (m.timestamp)",
                    "CREATE INDEX event_timestamp_index IF NOT EXISTS FOR (e:Event) ON (e.timestamp)",
                    "CREATE INDEX memory_type_index IF NOT EXISTS FOR (m:Memory) ON (m.memory_type)",
                    "CREATE INDEX event_type_index IF NOT EXISTS FOR (e:Event) ON (e.event_type)"
                ]
                
                for index in indexes:
                    try:
                        await session.run(index)
                    except Exception as e:
                        # 索引可能已存在，忽略错误
                        logger.debug(f"索引创建跳过: {e}")
                
                logger.info("Neo4j约束和索引创建完成")
                
        except Exception as e:
            logger.error(f"创建Neo4j约束和索引失败: {e}")
    
    @asynccontextmanager
    async def get_session(self) -> AsyncSession:
        """获取数据库会话的上下文管理器"""
        if not self.driver:
            raise RuntimeError("数据库未连接")
        
        session = self.driver.session()
        try:
            yield session
        finally:
            await session.close()
    
    async def execute_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """执行Cypher查询"""
        if not self.driver:
            raise RuntimeError("数据库未连接")
        
        for attempt in range(self._max_retry_attempts):
            try:
                async with self.get_session() as session:
                    result = await session.run(query, parameters or {})
                    records = await result.data()
                    return records
                    
            except (ServiceUnavailable, TransientError) as e:
                if attempt < self._max_retry_attempts - 1:
                    logger.warning(f"查询执行失败，尝试重试 ({attempt + 1}/{self._max_retry_attempts}): {e}")
                    await asyncio.sleep(self._retry_delay * (2 ** attempt))
                    continue
                else:
                    logger.error(f"查询执行失败，已达到最大重试次数: {e}")
                    raise
            except Exception as e:
                logger.error(f"执行Neo4j查询失败: {e}")
                raise
        
        return []

    async def execute_read_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """执行只读查询（兼容别名，委托 execute_query）"""
        return await self.execute_query(query, parameters)

    async def execute_write_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """执行写入查询"""
        if not self.driver:
            raise RuntimeError("数据库未连接")
        
        for attempt in range(self._max_retry_attempts):
            try:
                async with self.get_session() as session:
                    result = await session.run(query, parameters or {})
                    summary = await result.consume()
                    return {
                        'nodes_created': summary.counters.nodes_created,
                        'relationships_created': summary.counters.relationships_created,
                        'properties_set': summary.counters.properties_set
                    }
                    
            except (ServiceUnavailable, TransientError) as e:
                if attempt < self._max_retry_attempts - 1:
                    logger.warning(f"写入查询执行失败，尝试重试 ({attempt + 1}/{self._max_retry_attempts}): {e}")
                    await asyncio.sleep(self._retry_delay * (2 ** attempt))
                    continue
                else:
                    logger.error(f"写入查询执行失败，已达到最大重试次数: {e}")
                    raise
            except Exception as e:
                logger.error(f"执行Neo4j写入查询失败: {e}")
                raise
        
        return {'nodes_created': 0, 'relationships_created': 0, 'properties_set': 0}
    
    # 用户节点操作
    async def create_user_node(self, user_id: str, user_data: Dict[str, Any]) -> bool:
        """创建用户节点"""
        try:
            query = """
            MERGE (u:User {user_id: $user_id})
            SET u += $user_data
            SET u.created_at = datetime()
            SET u.updated_at = datetime()
            RETURN u
            """
            
            parameters = {
                'user_id': user_id,
                'user_data': user_data
            }
            
            await self.execute_write_query(query, parameters)
            logger.debug(f"用户节点已创建: {user_id}")
            return True
            
        except Exception as e:
            logger.error(f"创建用户节点失败: {e}")
            return False
    
    async def get_user_node(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户节点"""
        try:
            query = "MATCH (u:User {user_id: $user_id}) RETURN u"
            result = await self.execute_query(query, {'user_id': user_id})
            
            if result:
                return result[0]['u']
            return None
            
        except Exception as e:
            logger.error(f"获取用户节点失败: {e}")
            return None
    
    # 概念节点操作
    async def create_concept_node(self, concept_id: str, concept_data: Dict[str, Any]) -> bool:
        """创建概念节点"""
        try:
            query = """
            MERGE (c:Concept {concept_id: $concept_id})
            SET c += $concept_data
            SET c.created_at = datetime()
            SET c.updated_at = datetime()
            RETURN c
            """
            
            parameters = {
                'concept_id': concept_id,
                'concept_data': concept_data
            }
            
            await self.execute_write_query(query, parameters)
            logger.debug(f"概念节点已创建: {concept_id}")
            return True
            
        except Exception as e:
            logger.error(f"创建概念节点失败: {e}")
            return False
    
    async def find_concepts_by_name(self, name: str, limit: int = 10) -> List[Dict[str, Any]]:
        """根据名称查找概念（不区分大小写）"""
        try:
            query = """
            MATCH (c:Concept)
            WHERE toLower(c.name) CONTAINS toLower($name)
            RETURN c
            LIMIT $limit
            """
            
            result = await self.execute_query(query, {'name': name, 'limit': limit})
            return [record['c'] for record in result]
            
        except Exception as e:
            logger.error(f"查找概念失败: {e}")
            return []
    
    # 记忆节点操作
    async def create_memory_node(self, memory_id: str, memory_data: Dict[str, Any]) -> bool:
        """创建记忆节点"""
        try:
            query = """
            MERGE (m:Memory {memory_id: $memory_id})
            SET m += $memory_data
            SET m.created_at = datetime()
            SET m.updated_at = datetime()
            RETURN m
            """
            
            parameters = {
                'memory_id': memory_id,
                'memory_data': memory_data
            }
            
            await self.execute_write_query(query, parameters)
            logger.debug(f"记忆节点已创建: {memory_id}")
            return True
            
        except Exception as e:
            logger.error(f"创建记忆节点失败: {e}")
            return False
    
    # 关系操作
    async def create_relationship(self, from_node_id: str, from_node_type: str, 
                                to_node_id: str, to_node_type: str, 
                                relationship_type: str, properties: Optional[Dict[str, Any]] = None) -> bool:
        """创建关系"""
        try:
            query = f"""
            MATCH (from:{from_node_type} {{{from_node_type.lower()}_id: $from_node_id}})
            MATCH (to:{to_node_type} {{{to_node_type.lower()}_id: $to_node_id}})
            MERGE (from)-[r:{relationship_type}]->(to)
            SET r += $properties
            SET r.created_at = datetime()
            RETURN r
            """
            
            parameters = {
                'from_node_id': from_node_id,
                'to_node_id': to_node_id,
                'properties': properties or {}
            }
            
            await self.execute_write_query(query, parameters)
            logger.debug(f"关系已创建: {from_node_id} -[{relationship_type}]-> {to_node_id}")
            return True
            
        except Exception as e:
            logger.error(f"创建关系失败: {e}")
            return False
    
    async def get_relationships(self, node_id: str, node_type: str, 
                              relationship_type: Optional[str] = None, 
                              direction: str = 'both', limit: int = 50) -> List[Dict[str, Any]]:
        """获取节点的关系"""
        try:
            if direction == 'outgoing':
                direction_clause = f"-[r{':' + relationship_type if relationship_type else ''}]->"
            elif direction == 'incoming':
                direction_clause = f"<-[r{':' + relationship_type if relationship_type else ''}]-"
            else:  # both
                direction_clause = f"-[r{':' + relationship_type if relationship_type else ''}]-"
            
            query = f"""
            MATCH (n:{node_type} {{{node_type.lower()}_id: $node_id}}){direction_clause}(other)
            RETURN r, other
            LIMIT $limit
            """
            
            result = await self.execute_query(query, {'node_id': node_id, 'limit': limit})
            return result
            
        except Exception as e:
            logger.error(f"获取关系失败: {e}")
            return []
    
    # 路径查询
    async def find_shortest_path(self, from_node_id: str, from_node_type: str, 
                               to_node_id: str, to_node_type: str, 
                               max_depth: int = 5) -> Optional[List[Dict[str, Any]]]:
        """查找最短路径"""
        try:
            query = f"""
            MATCH (from:{from_node_type} {{{from_node_type.lower()}_id: $from_node_id}})
            MATCH (to:{to_node_type} {{{to_node_type.lower()}_id: $to_node_id}})
            MATCH path = shortestPath((from)-[*1..{max_depth}]-(to))
            RETURN path
            """
            
            result = await self.execute_query(query, {
                'from_node_id': from_node_id,
                'to_node_id': to_node_id
            })
            
            if result:
                return result[0]['path']
            return None
            
        except Exception as e:
            logger.error(f"查找最短路径失败: {e}")
            return None
    
    async def find_related_concepts(self, concept_id: str, depth: int = 2, limit: int = 20) -> List[Dict[str, Any]]:
        """查找相关概念"""
        try:
            query = f"""
            MATCH (source:Concept {{concept_id: $concept_id}})
            MATCH (source)-[*1..{depth}]-(related:Concept)
            WHERE related.concept_id <> $concept_id
            WITH source, related
            MATCH (source)-[]-(related)
            RETURN DISTINCT related.concept_id as concept_id,
                   related.name as name,
                   COUNT(*) as connection_strength
            ORDER BY connection_strength DESC
            LIMIT $limit
            """
            
            result = await self.execute_query(query, {
                'concept_id': concept_id,
                'limit': limit
            })
            
            return result
            
        except Exception as e:
            logger.error(f"查找相关概念失败: {e}")
            return []
    
    async def migrate_relationship_type(self, old_type: str = "RELATES_TO", new_type: str = "RELATED_TO", batch_size: int = 100) -> int:
        """批量迁移旧关系类型为新关系类型。

        Args:
            old_type (str): 旧的关系类型名称。
            new_type (str): 新的关系类型名称。
            batch_size (int): 每批处理的关系数量，避免单次事务过大。

        Returns:
            int: 成功迁移的关系数量。
        """
        try:
            total_migrated = 0
            while True:
                query = f"""
                MATCH ()-[r:{old_type}]->()
                WITH r LIMIT $batch_size
                CALL {{
                    WITH r
                    WITH startNode(r) AS from, endNode(r) AS to, r AS rel
                    MERGE (from)-[newRel:{new_type}]->(to)
                    SET newRel += rel{{.*}},
                        newRel.migrated_at = datetime()
                    DELETE rel
                    RETURN 1 AS migrated
                }}
                RETURN sum(migrated) AS migrated_count
                """
                result = await self.execute_write_query(query, {"batch_size": batch_size})
                if not result:
                    break
                migrated_in_batch = result[0].get("migrated_count", 0)
                total_migrated += migrated_in_batch or 0
                if migrated_in_batch < batch_size:
                    # 本批处理数量小于batch_size，说明已迁移完毕
                    break
            logger.info(f"迁移完成，共迁移 {total_migrated} 条关系类型 {old_type} -> {new_type}")
            return total_migrated
        except Exception as e:
            logger.error(f"迁移关系类型失败: {e}")
            return 0

    async def find_related_memories(self, query: str, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """根据查询词和用户查找相关记忆

        Args:
            query (str): 搜索关键词
            user_id (str): 用户ID，仅搜索该用户的记忆
            limit (int): 返回数量上限

        Returns:
            List[Dict[str, Any]]: 记忆ID及关联度分数列表
        """
        try:
            cypher = """
            MATCH (u:User {user_id: $user_id})<-[:BELONGS_TO]-(m:Memory)
            WHERE toLower(m.content) CONTAINS toLower($query)
               OR (exists(m.title) AND toLower(m.title) CONTAINS toLower($query))
               OR any(tag IN coalesce(m.tags, []) WHERE toLower(tag) CONTAINS toLower($query))
            WITH m, size((m)--()) AS degree
            RETURN m.memory_id AS memory_id,
                   degree AS relation_score
            ORDER BY relation_score DESC
            LIMIT $limit
            """

            records = await self.execute_query(cypher, {
                "query": query,
                "user_id": user_id,
                "limit": limit
            })
            return records
        except Exception as e:
            logger.error(f"查找相关记忆失败: {e}")
            return []

    async def find_related_knowledge(self, query: str, user_id: str, limit: int = 5) -> List[Dict[str, Any]]:
        """根据查询词和用户查找相关知识文档

        Args:
            query (str): 搜索关键词
            user_id (str): 用户ID，仅搜索该用户上传的知识
            limit (int): 返回数量上限

        Returns:
            List[Dict[str, Any]]: 文档ID及关联度分数列表
        """
        try:
            cypher = """
            MATCH (u:User {user_id: $user_id})<-[:OWNED_BY]-(d:Document)
            WHERE toLower(d.title) CONTAINS toLower($query)
               OR toLower(d.content) CONTAINS toLower($query)
            WITH d, size((d)--()) AS degree
            RETURN d.document_id AS document_id,
                   degree AS relevance
            ORDER BY relevance DESC
            LIMIT $limit
            """
            records = await self.execute_query(cypher, {
                "query": query,
                "user_id": user_id,
                "limit": limit
            })
            return records
        except Exception as e:
            logger.error(f"查找相关知识失败: {e}")
            return []

    # ... existing code ...